(*
 * PLL OPTIMIZATION CODE
 *
 * Example usage: Read in model and data, optimize PLL, and write out
 * optimized MN file.
 *
 * let opt modelfile datafile outfile = 
 *   let mn = Mn.load_auto modelfile in
 *   let data = Array.of_list (Data.input_example_list (open_in datafile)) in
 *   let (cache, w) = Mn.build_pll_cache mn data in
 *   let (errcode, pll) = Lbfgs.minimize (pll_val_and_grad cache) w in
 *   let mn' = pll_cache_to_mn cache w in
 *   write_auto outfile mn'
 *)
open Ext
open Printf

(* Keep track of how many violations of a conjunctive feature are present --
   zero, one, or more than one. *)
type satisfaction = Sat | Unsat of Mn.Factor.condition | NeverSat

exception Break

let string_of_sat = function
  | Sat -> "Sat"
  | Unsat(sense, var, value) -> 
    sprintf "Unsat: %s%d_%d" (if sense then "+" else "-") var value
  | NeverSat -> "NeverSat"

let compute_sat x f = 
  let result = ref Sat in
  try
    for i = 0 to Array.length f.Mn.Factor.cond - 1 do
      let (sense, var, value) = f.Mn.Factor.cond.(i) in
      if (x.(var) = value) <> sense then
        if !result = Sat then
          result := Unsat(sense, var, value)
        else
          (result := NeverSat;
          raise Break)
    done;
    !result
  with Break -> !result

(* Generate set of other values of var *)
let rec sat_set_rec var value i =
  if i < 0 then 
    []
  else if i = value then 
    sat_set_rec var value (i-1)
  else 
    (var, i) :: (sat_set_rec var value (i-1))
    
(* Generate list of var/value assignments that satisfy the condition *)
let sat_set schema (sense, var, value) =
  if sense then
    [(var, value)]
  else
    sat_set_rec var value (schema.(var) - 1)


(* Perform action inside, rather than generating list... *)
let sat_add schema featurelists fi (sense, var, value) =
  if sense then
    featurelists.(var).(value) <- fi :: featurelists.(var).(value)
  else
    for i = 0 to schema.(var) - 1 do
      if i <> value then
        featurelists.(var).(i) <- fi :: featurelists.(var).(i)
    done


(* Generate set of var/value assignments that satisfy this feature,
   assuming all other variables are held constant.  Excludes variables
   that do not appear in the feature. *)
let generate_sat_states schema x f =  
  match compute_sat x f with
  | Sat -> 
    (* dlogf "Sat\n"; *)
    (* STOP TODO -- this is slow! *)
    List.flatten (List.map (sat_set schema) (Array.to_list f.Mn.Factor.cond))
  | Unsat (sense, var, value) -> 
    (* dlogf "Unsat: (%d, %d)\n" var (if sense then value else (-value-1)); *)
    sat_set schema (sense, var, value)
  | NeverSat -> 
    (* dlogf "Never\n"; *) []

(* PLL debugging support *)
let string_of_sat (var, value) =
  sprintf "%d_%d" var value

let string_of_satl l =
  String.concat " " (List.map string_of_sat l)
  
(* Helper function for computing the gradient *)
let add_counts_to_gradient schema featurelists (wt, x) w g =
  let value = ref 0.0 in
  for j = 0 to Array.length schema - 1 do 
    let jfeatures = featurelists.(j) in
    let logdist = Array.create schema.(j) 0.0 in 
    (* Compute distribution over P(X_j) *)
    for k = 0 to schema.(j) - 1 do
      logdist.(k) <- 0.0;
      for l = 0 to Array.length jfeatures.(k) - 1 do
        logdist.(k) <- logdist.(k) +. w.(jfeatures.(k).(l))
      done;
    done;
    normalize_inplace_log logdist;
    (* Add to value of PLL *)
    value := !value +. wt *. logdist.(x.(j));
    (* For each relevant feature's gradient, we want to add the 
     * difference between the true counts and the expectation.
     * We do this by subtracting one from the true state probability
     * (and zero from all others). *)
    (* Reuse the logdist array, to save memory *)
    let dist = logdist in
    for k = 0 to schema.(j) - 1 do
      dist.(k) <- exp logdist.(k)
    done;
    (* let dist = Array.map exp logdist in *)
    dist.(x.(j)) <- dist.(x.(j)) -. 1.;
    (* Add to gradient of PLL *)
    for k = 0 to schema.(j) - 1 do
      for l = 0 to Array.length jfeatures.(k) - 1 do
        let wi = jfeatures.(k).(l) in
        g.(wi) <- g.(wi) -. wt *. dist.(k)
      done
    done
  done;
  !value

let add_counts_to_gradient_l schema featurelists (wt, x) w g =
  let value = ref 0.0 in
  for j = 0 to Array.length schema - 1 do 
    let jfeatures = featurelists.(j) in
    (* Compute distribution over P(X_j) *)
    let logdist = Array.create schema.(j) 0.0 in
    for k = 0 to schema.(j) - 1 do
      logdist.(k) <- List.sumf_map (fun wi -> w.(wi)) jfeatures.(k)
    done;
    normalize_inplace_log logdist;
    (* Add to value of PLL *)
    value := !value +. wt *. logdist.(x.(j));
    (* For each relevant feature's gradient, we want to add the 
     * difference between the true counts and the expectation.
     * We do this by subtracting one from the true state probability
     * (and zero from all others). *)
    let dist = Array.map exp logdist in
    dist.(x.(j)) <- dist.(x.(j)) -. 1.;
    (* Add to gradient of PLL *)
    for k = 0 to schema.(j) - 1 do
      List.iter (fun wi -> g.(wi) <- g.(wi) -. wt *. dist.(k)) jfeatures.(k)
    done
  done;
  !value
      
(*
 * PLL Gradient Method 1: Build up a full cache first.  Works well. 
 *)

(* Construct key statistics for PLL value and gradient computations *)
let build_pll_cache mn data = 
  Timer.start "build_pll_cache";
  let schema = mn.Mn.schema in
  let fa = Array.of_list (Mn.to_features mn) in
  let datafeatures = Array.create (Array.length data) [||] in
  for i = 0 to Array.length data - 1 do
    let (wt, x) = data.(i) in
    let featurelists = Array.map (fun d -> Array.create d []) schema in
    let append_f fi (var,value) =
      featurelists.(var).(value) <- fi :: featurelists.(var).(value) in
    let add_f fi f =
      let sl = generate_sat_states schema x f in
      List.iter (append_f fi) sl in
    Array.iteri add_f fa;
    datafeatures.(i) <- Array.map (Array.map Array.of_list) featurelists
  done;
  Timer.stop "build_pll_cache";
  let cache = (schema, fa, data, datafeatures) in
  let w = Array.map (fun f -> f.Mn.Factor.weight) fa in
  (cache, w)


(* Compute gradient of the PLL. *)
let pll_val_and_grad_cached cache w g =
  (* Clear gradient array *)
  Timer.start "pll_val_and_grad";
  assert(Array.length w = Array.length g);
  for l = 0 to Array.length g - 1 do
    g.(l) <- 0.0
  done;
  let (schema, fa, data, datafeatures) = cache in
  let value = ref 0.0 in
  for i = 0 to Array.length data - 1 do
    value := !value +. add_counts_to_gradient schema datafeatures.(i) 
      data.(i) w g
  done;
  Timer.stop "pll_val_and_grad";
  !value

(*
 * PLL Gradient Method 2: Smaller cache, with only a list of satisfied
 * and almost-satisfied features.  Half as fast as method 1, and not
 * much memory savings.  Not recommended.
 *)
let build_pll_medcache mn data = 
  Timer.start "build_pll_cache";
  let schema = mn.Mn.schema in
  let fa = Array.of_list (Mn.to_features mn) in
  let satf = Array.create (Array.length data) [||] in
  let unsatf = Array.create (Array.length data) [||] in
  for i = 0 to Array.length satf - 1 do
    let (wt, x) = data.(i) in
    let satl = ref [] in
    let unsatl = ref [] in
    let satcount = ref 0 in
    let unsatcount = ref 0 in
    for fi = 0 to Array.length fa - 1 do
      match compute_sat x fa.(fi) with
      | NeverSat -> ()
      | Sat -> 
        incr satcount;
        satl := fi :: !satl
      | Unsat(sense, var, value) -> 
        incr unsatcount;
        unsatl := (fi, sense, var, value) :: !unsatl
    done;
    satf.(i) <- Array.of_list !satl;
    unsatf.(i) <- Array.of_list !unsatl;
    (* nlogf "%d: %d (%d %d)\n" i (Array.length counts.(i)) !satcount !unsatcount; *)
    flush stdout
  done;
  Timer.stop "build_pll_cache";
  let cache = (schema, fa, data, (satf, unsatf)) in
  let w = Array.map (fun f -> f.Mn.Factor.weight) fa in
  (cache, w)

(* Compute gradient of the PLL. *)
let pll_val_and_grad_medcached cache w g =
  (* Clear gradient array *)
  Timer.start "pll_val_and_grad";
  assert(Array.length w = Array.length g);
  for l = 0 to Array.length g - 1 do
    g.(l) <- 0.0
  done;
  let (schema, fa, data, (satf, unsatf)) = cache in
  let value = ref 0.0 in
  for i = 0 to Array.length satf - 1 do
    (* Compute relevant features *)
    let featurelists = Array.map (fun d -> Array.create d []) schema in
    for j = 0 to Array.length satf.(i) - 1 do
      let fi = satf.(i).(j) in
      for k = 0 to Array.length fa.(fi).Mn.Factor.cond - 1 do
        let c = fa.(fi).Mn.Factor.cond.(k) in
        sat_add schema featurelists fi c 
      done
    done;
    for j = 0 to Array.length unsatf.(i) - 1 do
      let (fi, sense, var, value) = unsatf.(i).(j) in
      sat_add schema featurelists fi (sense, var, value) 
    done;
    (* let featurelists = Array.map (Array.map Array.of_list) featurelists in *)
    value := !value +. add_counts_to_gradient_l schema featurelists
      data.(i) w g
  done;
  Timer.stop "pll_val_and_grad";
  !value

(*
 * PLL Gradient Method 3: No cache at all
 *)
let build_pll_minicache mn data = 
  let schema = mn.Mn.schema in
  let fa = Array.of_list (Mn.to_features mn) in
  let cache = (schema, fa, data, ()) in 
  let w = Array.map (fun f -> f.Mn.Factor.weight) fa in
  (cache, w)

let pll_val_and_grad cache w g =
  let (schema, fa, data, _) = cache in
  (* Clear gradient array *)
  Timer.start "pll_val_and_grad";
  assert(Array.length w = Array.length g);
  for l = 0 to Array.length g - 1 do
    g.(l) <- 0.0
  done;
  let value = ref 0.0 in
  for i = 0 to Array.length data - 1 do
    (* Compute features satisfied by each variable *)
    let featurelists = Array.map (fun d -> Array.create d []) schema in
    let append_f fi (var,value) =
      featurelists.(var).(value) <- fi :: featurelists.(var).(value) in
    let (weight, x) = data.(i) in
    let add_f fi f =
      let sl = generate_sat_states schema x f in
      List.iter (append_f fi) sl in
    Array.iteri add_f fa;
    (* Update value and gradient *)
    let featurelists = Array.map (Array.map Array.of_list) featurelists in
    value := !value +. add_counts_to_gradient schema featurelists
      data.(i) w g
  done;
  Timer.stop "pll_val_and_grad";
  !value


(* Construct an MN from the PLL cache data and the weight vector. *)
let pll_cache_to_mn cache w =
  let (schema, fa, data, _) = cache in
  Array.iter2 (fun f wi -> f.Mn.Factor.weight <- wi) fa w;
  let factors = Array.map (fun f -> Mn.Factor.Feature f) fa in
  let mn = Mn.create schema factors in
  mn

let fa_pll schema fa x =
  let score = ref 0.0 in
  let logdist = Array.map (fun d -> Array.create d 0.0) schema in
  let add_weight wi (var,value) =
    (* dlogf "value=%d w=%f" value wi; *)
    logdist.(var).(value) <- logdist.(var).(value) +. wi in
  let add_f f =
    let sl = generate_sat_states schema x f in
    List.iter (add_weight f.Mn.Factor.weight) sl in
  Array.iter add_f fa;
  Array.iter normalize_inplace_log logdist;
  for var = 0 to Array.length schema - 1 do
    (* dlogf "logdist.(%d).(%d) = %f\n" var x.(var) logdist.(var).(x.(var)); *)
    score := !score +. logdist.(var).(x.(var))
  done;
  !score

(* Support for a C encoding of MNs. *)
type mn_t

external create_mn_helper : 
    int array -> (bool * int * int) array array -> float array -> mn_t 
      = "ocaml_create_mn"

let create_mn schema fa =
  flush stdout;
  let w = Array.map (fun f -> f.Mn.Factor.weight) fa in
  let conds = Array.map (fun f -> f.Mn.Factor.cond) fa in
  create_mn_helper schema conds w

(* Support for a C encoding of a dataset. *)
type data_t

external create_data :
    (float * int array) array -> data_t = "ocaml_create_data"

(* External libraries for computing PLL and its gradient slightly faster. *)
external pll_mn :
    mn_t -> int array -> float = "ocaml_c_pll_mn"
    
external pll_val_and_grad_ext :
    mn_t -> data_t -> float array -> float array -> float = "ocaml_c_pll_val_and_grad"

let build_pll_ccache mn data =
  let schema = mn.Mn.schema in
  let fa = Array.of_list (Mn.to_features mn) in
  let cmn = create_mn schema fa in
  let cdata = create_data data in
  let w = Array.map (fun f -> f.Mn.Factor.weight) fa in
  let cache = (schema, fa, data, (cmn, cdata)) in
  (cache, w)


let pll_val_and_grad_c cache w g =
  flush stdout;
  Timer.start "pll_val_and_grad";
  let (_, _, _, (mn, data)) = cache in
  let pll = pll_val_and_grad_ext mn data w g in
  Timer.stop "pll_val_and_grad";
  pll

(* Support for C encoding of external cache info *)
type cache_t

external create_cache_ext :
  mn_t -> data_t -> cache_t = "ocaml_create_cache"

(* external check_cache_ext :
  cache_t -> cache_t = "ocaml_c_check_cache" *)

let build_pll_ccache2 mn data =
  let schema = mn.Mn.schema in
  let fa = Array.of_list (Mn.to_features mn) in
  let cmn = create_mn schema fa in
  let cdata = create_data data in
  let ccache = create_cache_ext cmn cdata in
  (* ignore(check_cache_ext ccache); *)
  let w = Array.map (fun f -> f.Mn.Factor.weight) fa in
  let cache = (schema, fa, data, (cmn, cdata, ccache)) in
  (cache, w)

external pll_val_and_grad_cache_ext :
    mn_t -> data_t -> cache_t -> float array -> float array -> float 
      = "ocaml_c_pll_val_and_grad_cached"

let pll_val_and_grad_ccache cache w g =
  flush stdout;
  Timer.start "pll_val_and_grad";
  let (_, _, _, (cmn, cdata, ccache)) = cache in
  (* ignore(check_cache_ext ccache); *)
  let pll = pll_val_and_grad_cache_ext cmn cdata ccache w g in
  Timer.stop "pll_val_and_grad";
  pll
